import pdb
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from loss.ap_loss import APLoss

class PixelAPLoss(nn.Module):
    def __init__(self, sampler, nq=20):
        nn.Module.__init__(self)
        self.aploss = APLoss(nq, min=0, max=1, euc=False)
        self.name = 'pixAP'
        self.sampler = sampler

    def loss_from_ap(self, ap):
        return 1 - ap

    def forward(self,outputs, inputs, **kw):
        # subsample things
        scores, gt, msk = self.sampler(outputs,inputs)
        # scores, gt, msk, qconf = self.sampler(descriptors, None, aflow)

        # compute pixel-wise AP
        n = scores.size()[0]
        if n == 0: return 0
        scores, gt = scores.view(n, -1), gt.view(n, -1)
        ap = self.aploss(scores, gt).view(msk.shape)
        # print("qconf: ", qconf.shape, torch.max(qconf), torch.min(qconf), torch.median(qconf))

        pixel_loss = self.loss_from_ap(ap)

        loss = pixel_loss[msk].mean()
        return loss